import pandas as pd
import numpy as np
import json
import os
import csv
def school_disambiguation(df, col="rept_univ"):
    content_replace = {"清华大学": ["清华"], "北京大学": ["北京大学"], "上海交大": ["上海交通", "上交"],
                       "浙江大学": ["浙大", "浙江大学"], "复旦大学": ["复旦"],
                       "中国科学技术大学": ["中科大", "中国科学技术大学"], "中国人民大学": ["中国人大", "中国人民大学"],
                       "上海财经大学": ["上海财经"], "厦门大学": ["厦门大学"], "武汉大学": ["武大", "武汉大学"],
                       "中南财经政法大学": ["中南财经", "中南财经政法大学"], "中央财经大学": ["中央财经大学"],
                       "中山大学": ["中山大学"], "南京大学": ["南京大学"],
                       "中科院": ["中科院", "中国科学院"]}
    start_replace = {"北京大学": ["北大"]}
    for k, v in content_replace.items():
        for val in v:
            df[col] = df[col].apply(lambda x: k if val in str(x) else x)
    for k, v in start_replace.items():
        for val in v:
            df[col] = df[col].apply(lambda x: k if str(x).startswith(val) else x)
    return df


def dept_disambiguation(df):
    content_replace = {"会计专业": ["会计"], "法律专业": ["法律", "法学"], "经管专业": ["经管"], "经济专业": ["经济"], "金融专业": ["金融"],
                       "工商管理专业": ["工商管理"]}
    for k, v in content_replace.items():
        for val in v:
            df["user_dept"] = df["user_dept"].apply(lambda x: k if val in str(x) else x)
    return df

def complete_basic_info(df):
    # 把nan的学历都设为最高学历
    df.replace('None', np.nan, inplace=True)
    # max_year_degree = df.loc[df['user_deg'].notnull()].groupby('uid')['user_deg'].transform('last')
    # df['age'] = df['age'].fillna(max_age)
    max_year_degree = df[df['user_deg'].notna()].sort_values('rept_year').groupby('uid')['user_deg'].last()

    def fill_in_as_before(x):
        try:
            print(x)
            x.fillna(max_year_degree[x.name])
        except:
            pass # all nan
    df['user_deg'] = df.groupby('uid', group_keys=False)['user_deg'].apply(fill_in_as_before)
    

def school_drop(df):
    drop_content = ["中国注册会计师"]
    for val in drop_content:
        df = df[df["rept_univ"] != val]
    return df


def dept_drop(df):
    # df = df[df["user_dept"] is not None]
    new_df = df.__deepcopy__()
    drop_content = ["None"]
    for val in drop_content:
        new_df = new_df[new_df["user_dept"] != val]
    return new_df

def age_drop(df):
    new_df = df.__deepcopy__()
    drop_content = ["None"]
    for val in drop_content:
        new_df = new_df[new_df["user_brth"] != val]
    return new_df


def only_last_year(df):
    df = df.loc[df.groupby("uid")["rept_year"].idxmax()]
    return df


def read_basic():
    basic_df = pd.read_csv("../../executive_basic_info.csv")
    # 每个人只取最后一个职位

    basic_df["user_deg"] = basic_df["user_deg"].replace("中职", "中专及以下")
    basic_df["user_deg"] = basic_df["user_deg"].replace("初中", "中专及以下")
    basic_df["user_deg"] = basic_df["user_deg"].replace("专科", "大专")
    
    basic_df['age'] = basic_df['rept_year'] - basic_df['user_brth']

    def calc_salary(x):
        for i, (l, r, f) in enumerate(salary):
            if l <= x <= r:
                return f
    
    def calc_age(x):
        for i, (l, r) in enumerate(age):
            if l <= x <= r:
                return f"{l}-{r}岁"

    basic_df["salary"] = basic_df["salary"].apply(calc_salary)
    basic_df["age"] = basic_df["age"].apply(calc_age)
    # basic_df = basic_df.drop('user_deg', axis=1)

    return basic_df


def get_company(year=False):
    basic_df = pd.read_csv("../../executive_basic_info.csv")
    if year:
        groupby_cols = ["rept_year", "stock_id"]
    else:
        groupby_cols = ["stock_id"]
    tmp = basic_df.groupby(groupby_cols)["uid"].nunique()


def get_deg(year=False):
    basic_df = read_basic()
    if year:
        groupby_cols = ["rept_year", "user_deg"]
    else:
        groupby_cols = ["user_deg"]
    tmp = basic_df.groupby(groupby_cols)["uid"].nunique()
    print(tmp)


def get_institute(year=False):
    institute_df = pd.read_csv("../../executive_institute_staff.csv")
    institute_df = school_disambiguation(institute_df, "user_univ")
    institute_df = school_drop(institute_df)
    if year:
        groupby_cols = ["rept_year", "rept_univ"]
    else:
        institute_df = only_last_year(institute_df)
        groupby_cols = ["rept_univ"]
    tmp = institute_df.groupby(groupby_cols)["uid"].nunique()
    tmp = tmp.sort_values(ascending=False)


def get_year_dept(year=False):
    education_df = pd.read_csv("../../executive_education_info.csv")
    education_df = dept_disambiguation(education_df)
    education_df = dept_drop(education_df)
    if year:
        groupby_cols = ["rept_year", "user_dept"]
    else:
        education_df = only_last_year(education_df)
        groupby_cols = ["user_dept"]
    tmp = education_df.groupby(groupby_cols)["uid"].nunique()
    tmp = tmp.sort_values(ascending=False)
    breakpoint()


def get_posn():
    basic_df = pd.read_csv("../../executive_basic_info.csv")
    groupby_cols = ["user_posn"]
    tmp = basic_df.groupby(groupby_cols)["uid"].nunique()
    breakpoint()


def get_salary(year=False):
    basic_df = pd.read_csv("../../executive_basic_info.csv")
    if not year:
        basic_df = only_last_year(basic_df)
    if year:
        groupby_cols = ["rept_year", "salary"]
    else:
        groupby_cols = ["salary"]
    tmp = basic_df.groupby(groupby_cols)["uid"].nunique()
    breakpoint()


show_content = {}
salary = [[1, 30000, "0-3万"], [30001, 100000, "3万-10万"], [100001, 200000, "10万-20万"], [200001, 500000, "20万-50万"], [500001, 1000000, "50万-100万"], [1000001, 2000000, "100万-200万"],
          [2000001, 5000000, "200万-500万"], [5000001, 50000000, "500万-INF"]]
age = [[20, 30], [31, 40], [41, 50], [51, 60], [61, 70], [71, 100]]
show_content["institute"] = ["清华大学", "北京大学", "上海交大", "浙江大学", "复旦大学", "中国科学技术大学",
                             "中国人民大学", "上海财经大学", "厦门大学", "武汉大学", "中南财经政法大学",
                             "中央财经大学", "中山大学", "南京大学", "中科院"]
show_content["degree"] = ["博士", "硕士", "本科", "大专", "中专及以下"]
show_content["department"] = ["会计专业", "法律专业", "经济专业", "金融专业", "工商管理专业", "经济管理专业"]
show_content["salary"] = ["0-3万", "3万-10万", "10万-20万", "20万-50万", "50万-100万",
                          "100万-200万",
                          "200万-500万", "500万-INF"]
show_content["age"] = ["20-30岁", "31-40岁", "41-50岁", "51-60岁", "61-70岁",
                          "71-100岁"]

# map_content = {}
# map_content["institute"] = ["清华大学", "北京大学", "上海交大", "浙江大学", "复旦大学", "中国科学技术大学",
#                              "中国人民大学", "上海财经大学", "厦门大学", "武汉大学", "中南财经政法大学",
#                              "中央财经大学", "中山大学", "南京大学", "中科院"]
# map_content["degree"] = ["博士", "硕士", "本科", "大专", "中专及以下"]
# map_content["department"] = ["会计专业", "法律专业", "经济专业", "金融专业", "工商管理专业", "经济管理专业"]
# map_content["salary"] = ["0-3万", "3万-10万", "10万-20万", "20万-50万", "50万-100万",
#                           "100万-200万",
#                           "200万-500万", "500万-INF"]
# map_content["age"] = ["20-30岁", "31-40岁", "41-50岁", "51-60岁", "61-70岁",
#                           "71-80", "81-100岁"]


class PeopleInfo:
    def __init__(self):
        self.basic_df = read_basic()
        self.education_df_with_user_deg = pd.read_csv("../../executive_education_info.csv")
        self.education_df = self.education_df_with_user_deg.drop(columns=["user_deg"])
        # self.education_df = pd.read_csv("../../executive_education_info.csv")
        #del
        self.institute_df = pd.read_csv("../../executive_institute_staff.csv")
        with open("../sample-data/stock_id2info.json", 'r', encoding="utf-8") as f:
            self.stock_id2info = json.load(f)
        dept_disambiguation(self.education_df)
        school_disambiguation(self.education_df, "user_univ")
        school_disambiguation(self.institute_df, "rept_univ")
        self.merged_df = pd.merge(pd.merge(self.basic_df, self.education_df, on=['stock_id', 'uid', 'rept_year'], how='outer'), self.institute_df, on=['stock_id', 'uid', 'rept_year'], how='outer')

        # complete_basic_info(self.education_df) # basic info和edution
    def get_split_people_info(self):
        reserve_columns = ["uid", "user_univ", "user_deg", "user_dept", "salary", "age"]
        new_df = only_last_year(self.merged_df)
        new_df = new_df[reserve_columns]
        new_df = new_df.drop_duplicates()
        new_df = new_df.reset_index(drop=True)
        new_df = new_df.dropna()
        new_df.rename(columns={"user_univ": "institute", "user_deg": "degree", "user_dept": "department"}, inplace=True)
        cols = ["degree", "department", "institute", "salary", "age"]
        for idx, col in enumerate(cols):
            new_df = new_df[new_df[col].isin(show_content[col])]
        new_df.to_csv("split_people_info.csv", index=False)
    def get_people_info(self, uid=None, name=None):
        if name:
            uid = self.basic_df[self.basic_df["user_name"] == name]["uid"].values[0]

        # now we get uid
        name = self.basic_df[self.basic_df["uid"] == uid]["user_name"].values[0]
        user_birth = self.basic_df[self.basic_df["uid"] == uid]["user_brth"].values[0]

        basic_df = self.basic_df[(self.basic_df["user_name"] == name) & (self.basic_df["user_brth"] == user_birth)]
        # education_df = self.education_df[self.education_df["uid"] == uid]
        education_df = self.education_df_with_user_deg[self.education_df_with_user_deg["uid"] == uid]
        institute_df = self.institute_df[self.institute_df["uid"] == uid]

        # 然后以什么格式返回？
        personal_info = {"name": basic_df["user_name"].values[0]}
        personal_info["posn_line"] = []
        calc_dict = {}
        for index, row in basic_df.iterrows():
            company = row["stock_id"]
            if calc_dict.get(company) is None:
                calc_dict[company] = []
            calc_dict[company].append((row["rept_year"], row["user_posn"]))
        start_year = int(basic_df["rept_year"].min())
        end_year = int(basic_df["rept_year"].max())
        personal_info["start_year"] = start_year
        personal_info["end_year"] = end_year
        max_character_per_unit = 80 // (end_year-start_year+1)
        for company in calc_dict.keys():
            calc_dict[company].sort(key=lambda x: x[0])
            pre_year = 0
            pre_posn = "none"
            for idx, (year, user_posn) in enumerate(calc_dict[company]):
                if year == pre_year:
                    continue
                pre_year = year
                if user_posn == pre_posn:
                    display_value = ""
                else:
                    room = 1
                    for ii in range(idx+1, len(calc_dict[company])):
                        if calc_dict[company][ii][1] == user_posn:
                            room += 1
                        else:
                            break
                    max_character = max_character_per_unit * room 
                    end_idx = 0
                    cnter = 0
                    while cnter < max_character and end_idx < len(user_posn):
                        cnter = cnter+1 if user_posn[end_idx] != ',' else cnter+0.5
                        end_idx += 1
                    if end_idx == len(user_posn):
                        display_value = user_posn
                    else:
                        display_value = user_posn[:end_idx-1] + ".."
                        
                pre_posn = user_posn
                try:
                    display_company = self.stock_id2info[str(company)][0]
                except:
                    display_company = str(company)
                personal_info["posn_line"].append(
                    {"group": str(year), "variable": display_company, "value": user_posn, "display_value": display_value})
        
        # education info
        education_info = {"name": basic_df["user_name"].values[0]}
        education_info["posn_line"] = []
        calc_dict = {}
        for index, row in education_df.iterrows():
            school = row["user_univ"]
            if calc_dict.get(school) is None:
                calc_dict[school] = []
            tmp_dept = row["user_dept"] if not pd.isna(row["user_dept"]) else ""
            tmp_deg = row["user_deg"] if not pd.isna(row["user_deg"]) else ""
            calc_dict[school].append((row["rept_year"], tmp_dept+tmp_deg))
        try:
            start_year = int(education_df["rept_year"].min())
            end_year = int(education_df["rept_year"].max())
        except:
            start_year = int(basic_df["rept_year"].min())
            end_year = int(basic_df["rept_year"].max())
        education_info["start_year"] = start_year
        education_info["end_year"] = end_year
        max_character_per_unit = 80 // (end_year-start_year+1)
        for school in calc_dict.keys():
            calc_dict[school].sort(key=lambda x: x[0])
            pre_year = 0
            pre_posn = "none"
            for year, user_posn in calc_dict[school]:
                if year == pre_year:
                    continue
                pre_year = year
                if user_posn == pre_posn:
                    display_value = ""
                else:
                    room = 1
                    for ii in range(idx+1, len(calc_dict[school])):
                        if calc_dict[school][ii-1][0] == calc_dict[school][ii][0]:
                            continue
                        if calc_dict[school][ii][1] == user_posn:
                            room += 1
                        else:
                            break
                    max_character = max_character_per_unit * room 
                    end_idx = 0
                    cnter = 0
                    while cnter < max_character and end_idx < len(user_posn):
                        cnter = cnter+1 if user_posn[end_idx] != ',' else cnter+0.5
                        end_idx += 1
                    if end_idx == len(user_posn):
                        display_value = user_posn
                    else:
                        display_value = user_posn[:end_idx-1] + ".."
                        
                pre_posn = user_posn
                education_info["posn_line"].append(
                    {"group": str(year), "variable": school, "value": user_posn, "display_value": display_value})

        # institute info
        institute_info = {"name": basic_df["user_name"].values[0]}
        institute_info["posn_line"] = []
        calc_dict = {}
        for index, row in institute_df.iterrows():
            school = row["rept_univ"]
            if calc_dict.get(school) is None:
                calc_dict[school] = []
            tmp_dept = row["rept_dept"] if not pd.isna(row["rept_dept"]) else ""
            tmp_deg = row["univ_posn"] if not pd.isna(row["univ_posn"]) else ""
            calc_dict[school].append((row["rept_year"], tmp_dept+tmp_deg))
        try:
            start_year = int(institute_df["rept_year"].min())
            end_year = int(institute_df["rept_year"].max())
        except:
            start_year = int(basic_df["rept_year"].min())
            end_year = int(basic_df["rept_year"].max())
        institute_info["start_year"] = start_year
        institute_info["end_year"] = end_year
        max_character_per_unit = 80 // (end_year-start_year+1)
        for school in calc_dict.keys():
            calc_dict[school].sort(key=lambda x: x[0])
            pre_year = 0
            pre_posn = "none"
            for year, user_posn in calc_dict[school]:
                if year == pre_year:
                    continue
                pre_year = year
                if user_posn == pre_posn:
                    display_value = ""
                else:
                    room = 1
                    for ii in range(idx+1, len(calc_dict[school])):
                        if calc_dict[school][ii][1] == user_posn:
                            room += 1
                        else:
                            break
                    max_character = max_character_per_unit * room 
                    end_idx = 0
                    cnter = 0
                    while cnter < max_character and end_idx < len(user_posn):
                        cnter = cnter+1 if user_posn[end_idx] != ',' else cnter+0.5
                        end_idx += 1
                    if end_idx == len(user_posn):
                        display_value = user_posn
                    else:
                        display_value = user_posn[:end_idx-1] + ".."
                        
                pre_posn = user_posn
                institute_info["posn_line"].append(
                    {"group": str(year), "variable": school, "value": user_posn, "display_value": display_value})
        
        
        return_info = {"company": personal_info, "education": education_info, "institute": institute_info}
        # for index, row in basic_df.iterrows():
        print(return_info)
        return return_info


    def get_static(self, feature="salary,age", year=False):
        os.makedirs("../../cache", exist_ok=True)
        cache_name = "../../cache/static_" + feature + "_" + str(year)
        if os.path.exists(cache_name):
            return json.load(open(cache_name, 'r'))
        feature2col = {"company": "stock_id", "degree": "user_deg", "institute": "rept_univ", "department": "user_dept",
                       "salary": "salary", "age": "age"}
        # 如果feature由多个项目组成，必须是在一个table里的 # merge后了就不用了
        # table_basic = ["stock_id", "salary", "user_brth", "user_deg"]
        # table_dept = ["user_dept"]
        # table_univ = ["rept_univ"]
        features = feature.split(",")
        cols = []
        for feat in features:
            cols.append(feature2col[feat])
            
        query_df = self.merged_df
        if "user_dept" in cols:
            query_df = dept_drop(query_df)
        if "rept_univ" in cols:
            query_df = school_drop(query_df)
        if not year:
            query_df = only_last_year(query_df)
 
        if year:
            groupby_cols = ["rept_year", ] + cols
        else:
            groupby_cols = cols
        
        tmp = query_df.groupby(groupby_cols)["uid"].nunique()
        # if year:
        for idx, col in enumerate(cols):
            tmp = tmp[tmp.index.get_level_values(col).isin(show_content[features[idx]])]
        # else:
        #     tmp = tmp[show_content[feature]]
        tmp = tmp.to_json(force_ascii=False)
        tmp = json.loads(tmp)
        tmp = list(tmp.items())
        # keys = tmp[0][0].split(",")
        # keys = [x.strip("() '") for x in keys]
        def map2dict(x):
            if year:
                date, name = x[0].split(",")
                date = date[1:]
                name = name[2:-2]
                return {"date": str(date)+"-01-01", "name": name, "value": x[1]}
            else:
                keys = x[0].split(",")
                data = {}
                for idx, key in enumerate(keys):
                    data[features[idx]] = key.strip("() '")
                data["value"] = x[1]
                return data

        tmp = list(map(map2dict, tmp))
        
        if not year and len(features)==2:
            data = []
            key1_set = set([x[features[0]] for x in tmp])
            for key1 in key1_set:
                tmp_data = {"name": key1}
                for key2 in show_content[features[1]]:
                    tmp_data[key2] = 0
                for item in tmp:
                    if item[features[0]] == key1:
                        tmp_data[item[features[1]]] += item["value"]
                data.append(tmp_data)
            tmp = data
            tmp = {"data": tmp, "group": show_content[features[0]], "subname": show_content[features[1]]}
        json.dump(tmp, open(cache_name, 'w'), ensure_ascii=False)
        return tmp

    def get_posn(self):
        basic_df = pd.read_csv("../../executive_basic_info.csv")
        groupby_cols = ["user_posn"]
        tmp = basic_df.groupby(groupby_cols)["uid"].nunique()
        breakpoint()
        
        
        


if __name__ == "__main__":
    # basic_df = None
    # get_year_company(basic_df)
    # get_year_deg()
    # get_deg()
    # get_year_institute()
    # get_year_dept()
    # get_posn()
    # get_salary()
    # get_people_info(name="戚聿东")
    people_info = PeopleInfo()
    # tmp = people_info.get_static("age,salary", year=True)
    # tmp = people_info.get_static("salary,age", year=False)
    people_info.get_split_people_info()
    breakpoint()
    pass
